package com.dk.shared.executor;

import com.dk.shared.utils.UserAuditorHelper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.task.TaskRejectedException;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;

/**
 * @Description
 * @Author 王大凯
 * @Date 2024-05-09 9:48
 */
@Slf4j
public class CommonTaskExecutor  {

    public CommonTaskExecutor(ThreadPoolTaskExecutor executor) {
        this.executor = executor;
    }

    private ThreadPoolTaskExecutor executor;

    public void execute(Runnable task) {
        String token = UserAuditorHelper.retrieveCurrentToken().orElse(null);
        try {
            executor.execute(attachSecurityAndRequestContext(token,task));
        }
        catch (RejectedExecutionException ex) {
            throw new TaskRejectedException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }


    public Future<?> submit(Runnable task) {
        String token = UserAuditorHelper.retrieveCurrentToken().orElse(null);
        try {
            return executor.submit(attachSecurityAndRequestContext(token,task));
        }
        catch (RejectedExecutionException ex) {
            throw new TaskRejectedException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }


    public <T> Future<T> submit(Callable<T> task) {
        String token = UserAuditorHelper.retrieveCurrentToken().orElse(null);
        try {
            return executor.submit(attachSecurityAndRequestContext(token,task));
        }
        catch (RejectedExecutionException ex) {
            throw new TaskRejectedException("Executor [" + executor + "] did not accept task: " + task, ex);
        }
    }

    private  <T> Callable<T> attachSecurityAndRequestContext(String usingUserToken,Callable<T> task) {
        return () -> {
            tokenLog(usingUserToken);
            UserAuditorHelper.setCurrentToken(usingUserToken);
            T call = task.call();
            UserAuditorHelper.removeCurrentToken();
            return call;
        };
    }

    private Runnable attachSecurityAndRequestContext(String usingUserToken, Runnable runnable) {
        return () -> {
            tokenLog(usingUserToken);
            UserAuditorHelper.setCurrentToken(usingUserToken);
            runnable.run();
            UserAuditorHelper.removeCurrentToken();
        };
    }

    private void tokenLog(String usingUserToken){
        log.debug("Executor设置的 token：{}", usingUserToken);
    }

}
